#!/usr/bin/env python3
"""
服务管理脚本
用于启动、停止、重启和监控服务
"""

import os
import sys
import signal
import time
import psutil
import argparse
import subprocess
from pathlib import Path

class ServiceManager:
    def __init__(self):
        self.project_dir = Path(__file__).parent
        self.pid_file = self.project_dir / "logs" / "gunicorn.pid"
        self.log_dir = self.project_dir / "logs"
        
        # 确保日志目录存在
        self.log_dir.mkdir(exist_ok=True)
    
    def is_running(self):
        """检查服务是否运行"""
        if not self.pid_file.exists():
            return False
        
        try:
            with open(self.pid_file, 'r') as f:
                pid = int(f.read().strip())
            
            # 检查进程是否存在
            return psutil.pid_exists(pid)
        except (ValueError, FileNotFoundError):
            return False
    
    def get_pid(self):
        """获取服务PID"""
        if not self.pid_file.exists():
            return None
        
        try:
            with open(self.pid_file, 'r') as f:
                return int(f.read().strip())
        except (ValueError, FileNotFoundError):
            return None
    
    def start(self, workers=4, port=5000, daemon=True):
        """启动服务"""
        if self.is_running():
            print("❌ 服务已在运行中")
            return False
        
        print("🚀 启动服务...")
        
        # 检查环境配置
        if not (self.project_dir / ".env").exists():
            print("❌ .env文件不存在，请先配置环境变量")
            return False
        
        # 构建启动命令
        cmd = [
            sys.executable, "-m", "gunicorn",
            "-w", str(workers),
            "-b", f"0.0.0.0:{port}",
            "--timeout", "120",
            "--access-logfile", str(self.log_dir / "access.log"),
            "--error-logfile", str(self.log_dir / "error.log"),
            "--pid", str(self.pid_file),
            "app:create_app()"
        ]
        
        if daemon:
            cmd.append("--daemon")
        
        try:
            # 切换到项目目录
            os.chdir(self.project_dir)
            
            # 启动服务
            subprocess.run(cmd, check=True)
            
            if daemon:
                # 等待一下确保服务启动
                time.sleep(2)
                if self.is_running():
                    print("✅ 服务启动成功")
                    print(f"PID: {self.get_pid()}")
                    print(f"访问地址: http://localhost:{port}")
                    return True
                else:
                    print("❌ 服务启动失败")
                    return False
            else:
                print("✅ 服务已启动（前台模式）")
                return True
                
        except subprocess.CalledProcessError as e:
            print(f"❌ 启动失败: {e}")
            return False
    
    def stop(self):
        """停止服务"""
        if not self.is_running():
            print("❌ 服务未运行")
            return False
        
        pid = self.get_pid()
        if not pid:
            print("❌ 无法获取服务PID")
            return False
        
        print(f"🛑 停止服务 (PID: {pid})...")
        
        try:
            # 发送SIGTERM信号
            os.kill(pid, signal.SIGTERM)
            
            # 等待进程结束
            for _ in range(10):
                if not psutil.pid_exists(pid):
                    break
                time.sleep(1)
            
            # 如果还没结束，强制杀死
            if psutil.pid_exists(pid):
                print("强制停止服务...")
                os.kill(pid, signal.SIGKILL)
                time.sleep(1)
            
            # 清理PID文件
            if self.pid_file.exists():
                self.pid_file.unlink()
            
            print("✅ 服务已停止")
            return True
            
        except ProcessLookupError:
            print("✅ 服务已停止")
            if self.pid_file.exists():
                self.pid_file.unlink()
            return True
        except Exception as e:
            print(f"❌ 停止失败: {e}")
            return False
    
    def restart(self, workers=4, port=5000):
        """重启服务"""
        print("🔄 重启服务...")
        self.stop()
        time.sleep(2)
        return self.start(workers, port)
    
    def status(self):
        """查看服务状态"""
        print("📊 服务状态:")
        
        if self.is_running():
            pid = self.get_pid()
            try:
                process = psutil.Process(pid)
                print(f"✅ 服务运行中")
                print(f"PID: {pid}")
                print(f"启动时间: {time.ctime(process.create_time())}")
                print(f"CPU使用率: {process.cpu_percent():.1f}%")
                print(f"内存使用: {process.memory_info().rss / 1024 / 1024:.1f} MB")
                print(f"线程数: {process.num_threads()}")
                
                # 检查端口
                connections = process.connections()
                for conn in connections:
                    if conn.status == 'LISTEN':
                        print(f"监听端口: {conn.laddr.port}")
                
            except psutil.NoSuchProcess:
                print("❌ 进程不存在")
        else:
            print("❌ 服务未运行")
        
        # 检查日志文件
        print("\n📝 日志文件:")
        for log_file in ["access.log", "error.log", "app.log"]:
            log_path = self.log_dir / log_file
            if log_path.exists():
                size = log_path.stat().st_size / 1024 / 1024
                print(f"  {log_file}: {size:.1f} MB")
            else:
                print(f"  {log_file}: 不存在")
    
    def logs(self, log_type="error", lines=50):
        """查看日志"""
        log_file = self.log_dir / f"{log_type}.log"
        
        if not log_file.exists():
            print(f"❌ 日志文件不存在: {log_file}")
            return
        
        print(f"📝 {log_type}日志 (最后{lines}行):")
        print("-" * 50)
        
        try:
            # 使用tail命令或Python实现
            if os.name == 'posix':  # Linux/macOS
                subprocess.run(['tail', '-n', str(lines), str(log_file)])
            else:  # Windows
                with open(log_file, 'r', encoding='utf-8') as f:
                    lines_list = f.readlines()
                    for line in lines_list[-lines:]:
                        print(line.rstrip())
        except Exception as e:
            print(f"❌ 读取日志失败: {e}")
    
    def health_check(self):
        """健康检查"""
        import requests
        
        print("🔍 健康检查...")
        
        try:
            response = requests.get("http://localhost:5000/health", timeout=10)
            if response.status_code == 200:
                data = response.json()
                print("✅ 服务健康")
                print(f"状态: {data.get('status')}")
                print(f"版本: {data.get('version')}")
            else:
                print(f"❌ 健康检查失败: HTTP {response.status_code}")
        except requests.exceptions.RequestException as e:
            print(f"❌ 健康检查失败: {e}")

def main():
    parser = argparse.ArgumentParser(description="图生图API服务管理")
    parser.add_argument('action', choices=['start', 'stop', 'restart', 'status', 'logs', 'health'],
                       help='操作类型')
    parser.add_argument('-w', '--workers', type=int, default=4, help='工作进程数')
    parser.add_argument('-p', '--port', type=int, default=5000, help='端口号')
    parser.add_argument('--foreground', action='store_true', help='前台运行')
    parser.add_argument('--log-type', choices=['access', 'error', 'app'], default='error',
                       help='日志类型')
    parser.add_argument('--lines', type=int, default=50, help='显示日志行数')
    
    args = parser.parse_args()
    
    manager = ServiceManager()
    
    if args.action == 'start':
        manager.start(args.workers, args.port, not args.foreground)
    elif args.action == 'stop':
        manager.stop()
    elif args.action == 'restart':
        manager.restart(args.workers, args.port)
    elif args.action == 'status':
        manager.status()
    elif args.action == 'logs':
        manager.logs(args.log_type, args.lines)
    elif args.action == 'health':
        manager.health_check()

if __name__ == "__main__":
    main()
